import matplotlib.pyplot as plt
from numpy import ndarray

# 显示图像
def display_images(original:ndarray, watermarked:ndarray, extractedwatermark:ndarray, psnr:float, nc:float):
    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.title("Original Image")
    plt.imshow(original, cmap='gray')
    plt.axis('off')

    plt.subplot(1, 3, 2)
    plt.title("Watermarked Image")
    plt.imshow(watermarked, cmap='gray')
    plt.axis('off')
    
    # 在图像下方显示PSNR值 和 NC值
    plt.text(0.5, -0.15, f'PSNR: {psnr:.2f} dB', size=12, ha='center', transform=plt.gca().transAxes)
    plt.text(0.5, -0.07, f'nc: {nc:.2f} ', size=12, ha='center', transform=plt.gca().transAxes)

    plt.subplot(1, 3, 3)
    plt.title("Extracted Watermark")
    plt.imshow(extractedwatermark, cmap='gray')
    plt.axis('off')

    plt.show()

# 显示图像函数
def display_images_with_attack(original_image, watermarked_image, attacked_image, attack_type, extracted_watermark, psnr_value, nc_value):

    plt.figure(figsize=(12, 8))
    plt.subplot(2, 2, 1)
    plt.imshow(original_image, cmap='gray')
    plt.title("Original Image")
    plt.axis('off')

    plt.subplot(2, 2, 2)
    plt.imshow(watermarked_image, cmap='gray')
    plt.title(f'Watermarked Image\nPSNR: {psnr_value:.2f}')    
    plt.axis('off')

    plt.subplot(2, 2, 3)
    plt.imshow(attacked_image, cmap='gray')
    plt.title(f'attacked_image\ntype: {attack_type}')
    plt.axis('off')

    plt.subplot(2, 2, 4)
    plt.imshow(extracted_watermark, cmap='gray')
    plt.title(f'extracted_watermark\nNC: {nc_value:.2f}')
    plt.axis('off')

    plt.show()